// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/memory/allocation.h>

#include <pollux/common/memory/memory.h>

namespace kumo::pollux::memory {
    Allocation::~Allocation() {
        if (pool_ != nullptr) {
            pool_->freeNonContiguous(*this);
        }
        // NOTE: exception throw on object destruction will cause process crash.
        if ((numPages_ != 0) || !runs_.empty()) {
            POLLUX_FAIL("Bad Allocation state on destruction: {}", toString());
        }
    }

    void Allocation::append(uint8_t *address, MachinePageCount numPages) {
        POLLUX_CHECK(
            runs_.empty() || address != runs_.back().data(),
            "Appending a duplicate address into a PageRun");
        if (MELON_UNLIKELY(numPages > Allocation::PageRun::kMaxPagesInRun)) {
            POLLUX_MEM_ALLOC_ERROR(fmt::format(
                "The number of pages to append {} exceeds the PageRun limit {}",
                numPages,
                Allocation::PageRun::kMaxPagesInRun));
        }
        numPages_ += numPages;
        runs_.emplace_back(address, numPages);
    }

    void Allocation::appendMove(Allocation &other) {
        for (auto &run: other.runs_) {
            numPages_ += run.numPages();
            runs_.push_back(std::move(run));
        }
        other.runs_.clear();
        other.numPages_ = 0;
    }

    void Allocation::findRun(uint64_t offset, int32_t *index, int32_t *offsetInRun)
    const {
        uint64_t skipped = 0;
        for (int32_t i = 0; i < runs_.size(); ++i) {
            uint64_t size = AllocationTraits::pageBytes(runs_[i].numPages());
            if (offset - skipped < size) {
                *index = i;
                *offsetInRun = static_cast<int32_t>(offset - skipped);
                return;
            }
            skipped += size;
        }
        POLLUX_UNREACHABLE(
            "Seeking to an out of range offset {} in Allocation with {} pages and {} runs",
            offset,
            numPages_,
            runs_.size());
    }

    std::string Allocation::toString() const {
        return fmt::format(
            "Allocation[numPages:{}, numRuns:{}, pool:{}]",
            numPages_,
            runs_.size(),
            pool_ == nullptr ? "null" : "set");
    }

    ContiguousAllocation::~ContiguousAllocation() {
        if (pool_ != nullptr) {
            pool_->freeContiguous(*this);
            pool_ = nullptr;
        }
        // NOTE: exception throw on object destruction will cause process crash.
        if ((data_ != nullptr) || (size_ != 0)) {
            POLLUX_FAIL("Bad ContiguousAllocation state on destruction: {}", toString());
        }
    }

    void ContiguousAllocation::set(void *data, uint64_t size, uint64_t maxSize) {
        data_ = data;
        size_ = size;
        maxSize_ = maxSize != 0 ? maxSize : size;
        sanityCheck();
    }

    void ContiguousAllocation::grow(MachinePageCount increment) {
        pool_->growContiguous(increment, *this);
    }

    void ContiguousAllocation::clear() {
        pool_ = nullptr;
        set(nullptr, 0);
    }

    MachinePageCount ContiguousAllocation::numPages() const {
        return AllocationTraits::numPages(size_);
    }

    std::optional<melon::Range<char *> > ContiguousAllocation::hugePageRange() const {
        auto begin = reinterpret_cast<uintptr_t>(data_);
        auto roundedBegin = bits::roundUp(begin, AllocationTraits::kHugePageSize);
        auto roundedEnd = (begin + maxSize_) / AllocationTraits::kHugePageSize *
                          AllocationTraits::kHugePageSize;
        if (roundedEnd <= roundedBegin) {
            return std::nullopt;
        }
        return melon::Range<char *>(
            reinterpret_cast<char *>(roundedBegin), roundedEnd - roundedBegin);
    }

    std::string ContiguousAllocation::toString() const {
        return fmt::format(
            "ContiguousAllocation[data:{}, size:{}, maxSize: {}, pool:{}]",
            data_,
            size_,
            maxSize_,
            pool_ == nullptr ? "null" : "set");
    }
} // namespace kumo::pollux::memory
